import torch.nn.functional as F
import logging
from .softmax_loss import CrossEntropyLabelSmooth
from .center_loss import CenterLoss
from .triplet_loss import TripletLoss, WeightedTripletLoss
import torch


def make_loss(cfg, num_classes, feat_dim=2048):  # modified by gu
    # feat_dim = 2048
    logger = logging.getLogger("reid_baseline.train")
    logger.info("feature dimension: {}".format(feat_dim))
    if 'triplet' in cfg.MODEL.METRIC_LOSS_TYPE:
        if cfg.MODEL.WEIGHTED_TRIPLET:
            if cfg.MODEL.NO_MARGIN:
                triplet = WeightedTripletLoss()
                logger.info("using soft margin Weighted Triplet Loss for training")
            else:
                triplet = WeightedTripletLoss(cfg.SOLVER.MARGIN)  # triplet loss
                logger.info("using Weighted Triplet Loss for training with margin:{}".format(cfg.SOLVER.MARGIN))
        else:
            if cfg.MODEL.NO_MARGIN:
                triplet = TripletLoss()
                logger.info("using soft margin triplet loss for training")
            else:
                triplet = TripletLoss(cfg.SOLVER.MARGIN)  # triplet loss
                logger.info("using Triplet Loss for training with margin:{}".format(cfg.SOLVER.MARGIN))

    center_criterion = CenterLoss(num_classes=num_classes, feat_dim=feat_dim, use_gpu=True)  # center loss
    if cfg.MODEL.IF_LABELSMOOTH == 'on':
        id_loss_func = CrossEntropyLabelSmooth(num_classes=num_classes)
        logger.info("label smooth on, numclasses:{}".format(num_classes))
    else:
        id_loss_func = F.cross_entropy

    def loss_func(score, feat, target):
        if cfg.MODEL.METRIC_LOSS_TYPE == 'triplet_center':
            if cfg.MODEL.IF_LABELSMOOTH == 'on':
                return cfg.MODEL.ID_LOSS_WEIGHT * id_loss_func(score, target) + \
                       cfg.MODEL.TRIPLET_LOSS_WEIGHT * triplet(feat, target)[0] + \
                       cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target)
            else:
                return cfg.MODEL.ID_LOSS_WEIGHT * id_loss_func(score, target) + \
                       cfg.MODEL.TRIPLET_LOSS_WEIGHT * triplet(feat, target)[0] + \
                       cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target)
        elif cfg.MODEL.METRIC_LOSS_TYPE == 'center':
            if cfg.MODEL.IF_LABELSMOOTH == 'on':
                return cfg.MODEL.ID_LOSS_WEIGHT * id_loss_func(score, target) + \
                       cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target)
            else:
                return cfg.MODEL.ID_LOSS_WEIGHT * id_loss_func(score, target) + \
                       cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target)
        elif cfg.MODEL.METRIC_LOSS_TYPE == 'triplet':
            if cfg.MODEL.IF_LABELSMOOTH == 'on':
                return cfg.MODEL.ID_LOSS_WEIGHT * id_loss_func(score, target) + \
                       cfg.MODEL.TRIPLET_LOSS_WEIGHT * triplet(feat, target)[0]
            else:
                return cfg.MODEL.ID_LOSS_WEIGHT * id_loss_func(score, target) + \
                       cfg.MODEL.TRIPLET_LOSS_WEIGHT * triplet(feat, target)[0]
        else:
            print('unexpected loss type')

    return loss_func, center_criterion


def make_loss_RMGL(cfg, num_classes, feat_dim=2048, USE_RMGL_loss=False):  # modified by gu
    # feat_dim = 2048
    logger = logging.getLogger("reid_baseline.train")
    logger.info("feature dimension: {}".format(feat_dim))
    if 'triplet' in cfg.MODEL.METRIC_LOSS_TYPE:
        if cfg.MODEL.WEIGHTED_TRIPLET:
            if cfg.MODEL.NO_MARGIN:
                triplet = WeightedTripletLoss()
                logger.info("using soft margin Weighted Triplet Loss for training")
            else:
                triplet = WeightedTripletLoss(cfg.SOLVER.MARGIN)  # triplet loss
                logger.info("using Weighted Triplet Loss for training with margin:{}".format(cfg.SOLVER.MARGIN))
        else:
            if cfg.MODEL.NO_MARGIN:
                triplet = TripletLoss()
                logger.info("using soft margin triplet loss for training")
            else:
                triplet = TripletLoss(cfg.SOLVER.MARGIN)  # triplet loss
                logger.info("using Triplet Loss for training with margin:{}".format(cfg.SOLVER.MARGIN))

    center_criterion = CenterLoss(num_classes=num_classes, feat_dim=feat_dim, use_gpu=True)  # center loss
    center_criterion_branch2 = CenterLoss(num_classes=num_classes, feat_dim=feat_dim * 5,
                                          use_gpu=True)  # branch2 center loss
    center_criterion_branch3 = CenterLoss(num_classes=num_classes, feat_dim=feat_dim * 7,
                                          use_gpu=True)  # branch3 center loss
    if cfg.MODEL.IF_LABELSMOOTH == 'on':
        id_loss_func = CrossEntropyLabelSmooth(num_classes=num_classes)
        logger.info("label smooth on, numclasses:{}".format(num_classes))
    else:
        id_loss_func = F.cross_entropy

    def id_loss_func_weight(score, target):
        return cfg.MODEL.ID_LOSS_WEIGHT * id_loss_func(score, target)

    def triplet_center_fun_weight(feat, target, branch="normal"):
        if branch == "branch2":
            return cfg.MODEL.TRIPLET_LOSS_WEIGHT * triplet(feat, target)[0] + \
                   cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion_branch2(feat, target)
        elif branch == "branch3":
            return cfg.MODEL.TRIPLET_LOSS_WEIGHT * triplet(feat, target)[0] + \
                   cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion_branch3(feat, target)
        else:
            return cfg.MODEL.TRIPLET_LOSS_WEIGHT * triplet(feat, target)[0] + \
                   cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target)

    def loss_func(score_feat, target):
        if not USE_RMGL_loss:
            "全部使用triplet id center loss"
            for index_score_feat in range(len(score_feat) // 2):
                score = score_feat[2 * index_score_feat]
                feat = score_feat[2 * index_score_feat + 1]
                if index_score_feat == 0:  # ID+triplet+center
                    loss = id_loss_func_weight(score, target) + triplet_center_fun_weight(feat, target)
                else:
                    loss += id_loss_func_weight(score, target) + triplet_center_fun_weight(feat, target)
        else:
            "RMGL分branch计算loss"
            branch2_feat = []  # branch2的所有feat拼接
            branch3_feat = []
            for index_score_feat in range(len(score_feat) // 2):
                score = score_feat[2 * index_score_feat]
                feat = score_feat[2 * index_score_feat + 1]
                if index_score_feat == 0:
                    id_loss = id_loss_func_weight(score, target)
                else:
                    id_loss += id_loss_func_weight(score, target)

                if index_score_feat < 5:
                    branch2_feat.append(feat)
                else:
                    branch3_feat.append(feat)
            branch2_feat = torch.cat(tuple(branch2_feat), dim=1)
            branch3_feat = torch.cat(tuple(branch3_feat), dim=1)
            loss = id_loss + \
                   triplet_center_fun_weight(branch2_feat, target, branch="branch2") + \
                   triplet_center_fun_weight(branch3_feat, target, branch="branch3")
        # loss /= (index_score_feat + 1)
        return loss

    if USE_RMGL_loss:
        return loss_func, [center_criterion_branch2, center_criterion_branch3]
    else:
        return loss_func, [center_criterion]
